# Custom utility nodes for ComfyUI workflows, for saving and loading presets, as well as simplifying various actions.
import os
import json
from comfy.samplers import KSAMPLER_NAMES, SCHEDULER_NAMES

scheduler_mapping = {
    "FlowMatchEulerDiscreteScheduler" : "beta",
    "DPMSolverMultistepScheduler" : "karras",
    "DDIMScheduler" : "ddim_uniform",
}
presets_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "presets")
model_json = {}
prompt_json = {}
model_names = []
prompt_names = []
tasks = ['image_generation', 'video_generation']

try:
    with open(f"{presets_dir}/model_presets.json", "r") as f:
        model_json = json.load(f)['models']
        model_names = [j['name'] for j in model_json[tasks[0]]]
    with open(f"{presets_dir}/prompt_presets.json", "r") as f:
        prompt_json = json.load(f)['prompt_examples']
        prompt_names = list(prompt_json[tasks[0]]['positive_prompts'].keys())
    print("Loaded presets")
except Exception as e:
    raise Exception(f"Could not load presets: {e}")

# Any Type trick
class AnyType(str):
    def __ne__(self, __value: object) -> bool:
        return False

any_typ = AnyType("*")

class ModelPresetsNode:
    @classmethod
    def INPUT_TYPES(cls):
        global tasks, model_names, prompt_names
        return {
            "required": {
                "task": (tasks, {"default": tasks[0]}),
            },
            "optional": {
                "model_name": (model_names, {"default":model_names[0]}),
                "prompt_type":(prompt_names, {"default":prompt_names[0]})
            }
        }
    CATEGORY = "octo_presets"
    RETURN_TYPES = (
        "STRING",
        "STRING",
        "INT",
        "FLOAT",
        "INT",
        "INT",
        "STRING",
        "STRING",
    )    
    RETURN_NAMES = (
        "scheduler",
        "sampling_method",
        "steps",
        "cfg",
        "width",
        "height", 
        "positive_prompt_example",
        "negative_prompt_example",
    )
    FUNCTION = "get_presets"

    def get_presets(self, task, model_name, prompt_type):
        global model_json
        if model_name not in model_json[task]:
            global model_names
            # Task has been updated, update model_list
            model_names = [j['name']for j in model_json[task]]
            model_name = model_names[0]

        for pre in model_json[task]:
            if pre['name'] == model_name:
                model_presets = pre['recommended_settings']

        prompt_presets = prompt_json[task]
        # default values
        scheduler = [scheduler_mapping.get(model_presets['scheduler'], "simple")]
        sampler = [model_presets["sampling_method"]]
        steps = model_presets["steps"]
        cfg = model_presets["cfg_scale"]
        width = model_presets.get("width", 512)
        height = model_presets.get("height", 512)
        try:
            positive_prompt_example = prompt_presets.get("positive_prompts").get(prompt_type, "universal")
        except:
            positive_prompt_example = "universal"
        try:
            negative_prompt_example = prompt_presets.get("negative_prompts").get(prompt_type, "universal")
        except:
            negative_prompt_example = "universal"
        return (scheduler, sampler, steps, cfg, width, height, positive_prompt_example, negative_prompt_example)
        

class LoadExperimentPresetsJson:
    @classmethod
    def INPUT_TYPES(cls):
        paths = [f for f in os.listdir(presets_dir) if f.endswith(".json")]
        return {
            "required": {
                "json_file": ("STRING", paths, {"default":paths[0]}),
            },
            "optional": {
                "index": ("INT", {"default":0, "min":0}),
                "increment_index":("BOOLEAN", {"default":True}),
                "seed":("INT", {"default":4206980085}),
                "steps": ("INT", {"default":20, "min":1, "step":1}),
                "cfg": ("FLOAT", {"default":7.0, "step":0.1}),
                "width": ("INT", {"default":512, "min":64}),
                "height": ("INT", {"default":512, "min":64}),
                "positive_prompt": ("STRING", {"default":"high quality, best quality, 8k, ultra-detailed"}),
                "negative_prompt": ("STRING", {"default":"low quality, worst quality, blurry, fuzzy"}),
                "sampler": (KSAMPLER_NAMES, {"default":KSAMPLER_NAMES[0]}),
                "scheduler": (SCHEDULER_NAMES, {"default":SCHEDULER_NAMES[0]}),
                "filename_prefix": ("STRING", {"default":"ComfyUI"})
            }
        }
    CATEGORY = "octo_presets"
    RETURN_TYPES = (
        "INT",
        "BOOLEAN",
        "INT",
        "INT",
        "FLOAT"
        "INT",
        "INT",
        "STRING",
        "STRING",
        "STRING",
        "STRING",
        "STRING",
        "STRING"
    )    
    RETURN_NAMES = (
        "index",
        "increment_index",
        "seed",
        "steps",
        "cfg",
        "width",
        "height", 
        "positive_prompt",
        "negative_prompt",
        "sampler",
        "scheduler",
        "filename_prefix",
        "new_prefix"
    )
    FUNCTION = "get_presets"

    def get_presets(json_file, index, increment_index, seed, steps, cfg, width, height,
                    positive_prompt, negative_prompt, sampler, scheduler, filename_prefix):
        global presets_dir
        with open(f"{presets_dir}/{json_file}", "r") as f:
            exps = json.load(f)
            size = len(exps)
            exp = exps[index]
            if increment_index :
                index = (index + 1)%size
        seed = exp.get('seed', seed)
        steps = exp.get('steps', steps)
        cfg = exp.get('cfg', cfg)
        width = exp.get('width', width)
        height = exp.get('height', height)
        positive_prompt = exp.get("positive_prompt", positive_prompt)
        negative_prompt = exp.get("negative_prompt", negative_prompt)
        sampler = exp.get("sampler", sampler)
        scheduler = exp.get("scheduler", scheduler)
        filename_prefix = exp.get("filename_prefix", filename_prefix)

        new_prefix = f"exp_steps{steps}_cfg{cfg}_w{width}_h{height}_{sampler}_{scheduler}"
        return (index, increment_index, steps, cfg, width, height,
                    positive_prompt, negative_prompt, sampler, scheduler, filename_prefix, new_prefix)

class SaveExperimentPresetsJson:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "json_file": ("STRING", {"default":"default.json"}),
            },
            "optional": {
                "seed": ("INT", {"default":-1}),
                "steps": ("INT", {"default":-1}),
                "cfg": ("FLOAT", {"default":-1}),
                "width": ("INT", {"default":-1}),
                "height": ("INT", {"default":-1}),
                "positive_prompt": ("STRING", {"default":""}),
                "negative_prompt": ("STRING", {"default":""}),
                "sampler": ("STRING", {"default":""}),
                "scheduler": ("STRING", {"default":""}),
                "filename_prefix": ("STRING", {"default":""})
            }
        }
    CATEGORY = "octo_presets"
    FUNCTION = "save_presets"
    RETURN_TYPES = ()

    def save_presets(json_file, seed, steps, cfg, width, height,
                    positive_prompt, negative_prompt, sampler, scheduler, filename_prefix):
        exp = {}
        if seed is not None and seed >=0:
            exp['seed'] = seed
        if steps is not None and steps>0:
            exp["steps"] = steps
        if cfg is not None and cfg>0:
            exp["cfg"] = cfg
        if width is not None and width > 0:
            exp["width"] = width
        if height is not None and height > 0:
            exp["height"] = height
        
        if positive_prompt is not None and positive_prompt != "":
            exp["positive_prompt"] = positive_prompt
        if negative_prompt is not None and negative_prompt != "":
            exp["negative_prompt"] = negative_prompt
        if filename_prefix is not None and filename_prefix != "":
            exp['filename_prefix'] = filename_prefix

        if sampler is not None and sampler in KSAMPLER_NAMES:
            exp["sampler"] = sampler
        if scheduler is not None and scheduler in SCHEDULER_NAMES:
            exp['scheduler'] = scheduler
        try:
            with open(f"{presets_dir}/{json_file}", "r") as f:
                exps = json.load(f)
        except:
            exps = []
        exps.append(exp)
        with open(f"{presets_dir}/{json_file}", "w") as f:
                json.dump(exps, f, indent=4)
                print(f"{presets_dir}/{json_file}")
        return ()

class AnytypeToStringNode:
    """
    Simple node that converts combo output to string
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "combo_input": (any_typ, {"default": ""}), 
            }
        }
    
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text_output",)
    FUNCTION = "convert_to_string"
    CATEGORY = "utilities"
    
    def convert_to_string(self, combo_input):
        """Convert combo selection to string"""
        return (str(combo_input),)


NODE_CLASS_MAPPINGS = {
    "Model Presets" : ModelPresetsNode,
    "Load Experiment Presets Json": LoadExperimentPresetsJson,
    "Save Experiment Presets Json": SaveExperimentPresetsJson,
    "Anything into string": AnytypeToStringNode
}

# Optionally, you can rename the node in the `NODE_DISPLAY_NAME_MAPPINGS` dictionary.
NODE_DISPLAY_NAME_MAPPINGS = {
    "Model Presets": "Model Presets",
    "Load Experiment Presets Json": "Load Experiment Presets from Json",
    "Save Experiment Presets Json": "Save Experiment Presets to Json",
    "Anything into string": "Anything into string"
}